package com.itheima.ai.controller;
import com.itheima.ai.entity.vo.Result;
import com.itheima.ai.repository.ChatHistoryRepository;
import com.itheima.ai.repository.FileRepository;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.QuestionAnswerAdvisor;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.ExtractedTextFormatter;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.reader.pdf.config.PdfDocumentReaderConfig;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.core.io.Resource;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.net.URLEncoder;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Objects;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;

@Slf4j
@RequiredArgsConstructor
@RestController
@RequestMapping("/ai/pdf")
public class PdfController {
private final FileRepository fileRepository;
private final VectorStore vectorStore;
/**
* ⽂件上传
*/
private final ChatClient pdfChatClient;
    private final ChatHistoryRepository chatHistroyRepository;
    @RequestMapping (value = "/chat",produces = "text/html;charset=UTF-8")
    public Flux<String> chat(String prompt, String chatId) {
        Resource file = fileRepository.getFile(chatId);
        if(!file.exists()){
            throw new RuntimeException("⽂件不存在！");
        }
        chatHistroyRepository.save("pdf",chatId);
        return pdfChatClient.prompt()
                .user(prompt)
                .advisors(a->a.param(CHAT_MEMORY_CONVERSATION_ID_KEY,chatId))
                .advisors(a->a.param(QuestionAnswerAdvisor.FILTER_EXPRESSION, "file_name == '"+file.getFilename()+"'"))
                .stream()
                .content();
    }
@RequestMapping("/upload/{chatId}")
public Result uploadPdf(@PathVariable String chatId, @RequestParam("file")
MultipartFile file) {
try {
// 1. 校验⽂件是否为PDF格式
if (!Objects.equals(file.getContentType(), "application/pdf")) {
return Result.fail("只能上传PDF⽂件！");
}
// 2.保存⽂件
boolean success = fileRepository.save(chatId, file.getResource());
if(! success) {
return Result.fail("保存⽂件失败！");
}
this.writeToVectorStore(file.getResource());
return Result.ok();
} catch (Exception e) {
log.error("Failed to upload PDF.", e);
return Result.fail("上传⽂件失败！");
}
}
/**
* ⽂件下载
*/
@GetMapping("/file/{chatId}")
public ResponseEntity<Resource> download(@PathVariable("chatId") String
chatId) throws IOException {
// 1.读取⽂件
Resource resource = fileRepository.getFile(chatId);
if (!resource.exists()) {
return ResponseEntity.notFound().build();
}
// 2.⽂件名编码，写⼊响应头
String filename =
URLEncoder.encode(Objects.requireNonNull(resource.getFilename()),
StandardCharsets.UTF_8);
// 3.返回⽂件
return ResponseEntity.ok()
.contentType(MediaType.APPLICATION_OCTET_STREAM)
.header("Content-Disposition", "attachment; filename=\"" +
filename + "\"")
.body(resource);
}
private void writeToVectorStore(Resource resource) {
// 1.创建PDF的读取器
PagePdfDocumentReader reader = new PagePdfDocumentReader(
resource, // ⽂件源
PdfDocumentReaderConfig.builder()
.withPageExtractedTextFormatter(ExtractedTextFormatter.defaults())
.withPagesPerDocument(1)
        .build()
);
    // 2.读取PDF文档，拆分为Document
    List<Document> documents = reader.read();
    // 3.写入向量库
    vectorStore.add(documents);
}
}